#! /usr/bin/env python
# -*- coding: utf-8 -*-

# Copyright (c) Huoty, All rights reserved
# Author: Huoty <sudohuoty@163.com>
# CreateTime: 2020-01-18 11:49:29

"""二叉树"""


class Node(object):

    def __init__(self, data, lchild=None, dchild=None):
        self.data = data
        self.lchild = None
        self.rchild = None


class BTree(object):
    """Binary Tree"""

    def __init__(self, root=None):
        self.root = root

    def __str__(self):
        """以逆时针旋转 90 度的视角打印树的结构"""
        def recurse(node, level):
            text = ""
            if node is not None:
                text += recurse(node.rchild, level + 1)
                text += "| " * level
                text += str(node.data) + "\n"
                text += recurse(node.lchild, level + 1)
            return text

        return recurse(self.root, 0).strip()

    def is_empty(self):
        return self.root is None

    def add(self, data):
        """添加节点"""
        node = Node(data)

        # 如果树为空，则设置根节点
        if self.is_empty():
            self.root = node
            return

        queue = []  # 用队列来记录节点链路
        queue.append(self.root)  # 从根开始遍历

        # 进行层次遍历
        while queue:
            # 从队列中取出节点
            probe = queue.pop(0)        # 每次取队头元素
            if probe.lchild is None:    # 如果左子树为空则添加为左节点
                probe.lchild = node
                return                  # 添加完成之后返回
            elif probe.rchild is None:  # 如果右子树为空则添加为右节点
                probe.rchild = node
                return                  # 添加完成之后返回
            else:
                # 如果左右子树都不为空，则加入队列继续向下遍历
                queue.append(probe.lchild)
                queue.append(probe.rchild)

    def add_as_bst(self, data):
        """插入为二叉搜索树"""
        def recurse(node):
            if data < node.data:
                if node.lchild is None:
                    node.lchild = Node(data)
                else:
                    recurse(node.lchild)
            elif node.rchild is None:
                node.rchild = Node(data)
            else:
                recurse(node.rchild)

        if self.is_empty():
            self.root = Node(data)
        else:
            recurse(self.root)

    def preorder(self):
        """前序遍历

        先访问根节点，然后再以类似的方式分别访问左子树和右子树
        （根，左，右；也可以是 根，右，左；所以遍历顺序与访问 根 的顺序有关）
        """
        def preorder_generator(root):
            if root is not None:
                yield root.data
                for data in preorder_generator(root.lchild):
                    yield data
                # yield from preorder_generator(root.lchild)  # Python > 3.3
                for data in preorder_generator(root.rchild):
                    yield data
                # yield from preorder_generator(root.rchild)  # Python > 3.3

        return preorder_generator(self.root)

    def preorder_stack(self):
        """前序遍历栈实现"""
        if self.root is not None:
            stack = []
            node = self.root
            while node or stack:
                # 从根节点开始，一直查找其左子树
                while node:
                    yield node.data
                    stack.append(node)
                    node = node.lchild
                # 到此，当前节点的左子树已经访问完毕，接着开始访问右子树
                node = stack.pop()
                node = node.rchild

    def inorder(self):
        """中序遍历

        先访问左节点，再访问根节点，最后访问右节点（左，根，右；或 右，根，左）
        """
        def inorder_generator(root):
            if root is not None:
                for data in inorder_generator(root.lchild):
                    yield data
                yield root.data
                for data in inorder_generator(root.rchild):
                    yield data

        return inorder_generator(self.root)

    def inorder_stack(self):
        """中序遍历栈实现"""
        if self.root is not None:
            stack = []
            node = self.root
            while node or stack:
                # 从根节点开始，一直查找其左子树
                while node:
                    stack.append(node)
                    node = node.lchild
                # 到此，当前节点的左子树已经访问完毕，开始访问右子树
                node = stack.pop()
                yield node.data
                node = node.rchild

    def postorder(self):
        """后续遍历

        先访问左节点，再访问右节点，最后根节点（左，右，根；或 右，左，根）
        """
        def postorder_generator(root):
            if root is not None:
                for data in postorder_generator(root.lchild):
                    yield data
                for data in postorder_generator(root.rchild):
                    yield data
                yield root.data

        return postorder_generator(self.root)

    def postorder_stack(self):
        """后续遍历栈实现"""
        if self.root is not None:
            stack1 = []
            stack2 = []
            node = self.root
            stack1.append(node)
            # 先按后序遍历的逆序压入到栈 stack2 中
            while stack1:
                node = stack1.pop()
                if node.lchild is not None:
                    stack1.append(node.lchild)
                if node.rchild is not None:
                    stack1.append(node.rchild)
                stack2.append(node)
            while stack2:
                yield stack2.pop().data

    def levelorder(self):
        """层序遍历（广度优先遍历）

        即按照层次访问，通常用队列来做
        访问根，访问子女，再访问子女的子女（越往后的层次越低，两个子女的级别相同）
        """
        if self.root is not None:
            queue = [self.root]
            while queue:
                node = queue.pop(0)  # pop node
                yield node.data
                # push next level
                if node.lchild is not None:
                    queue.append(node.lchild)
                if node.rchild is not None:
                    queue.append(node.rchild)


if __name__ == "__main__":
    btree = BTree()
    btree.add(1)
    btree.add(2)
    btree.add(3)
    btree.add(4)
    btree.add(5)
    btree.add(6)
    btree.add(7)

    print(btree)
    print("preorder:", list(btree.preorder()))
    print("preorder_stack:", list(btree.preorder_stack()))
    print("inorder:", list(btree.inorder()))
    print("inorder_stack:", list(btree.inorder_stack()))
    print("postorder:", list(btree.postorder()))
    print("postorder_stack:", list(btree.postorder_stack()))
    print("levelorder:", list(btree.levelorder()))

    print()
    bst = BTree()
    for data in [10, 5, 17, 3, 14, 9, 16]:
        bst.add_as_bst(data)
    print(bst)
    print("preorder:", list(bst.preorder()))
    print("inorder:", list(bst.inorder()))
